{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "heatmap   8 typhoons  RI-distance "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import os,glob,math\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.lines as mlines\n",
    "import matplotlib.dates as mdates\n",
    "import cartopy.crs as ccrs\n",
    "from cartopy.mpl.ticker import LongitudeFormatter,LatitudeFormatter\n",
    "import cartopy.feature as cfeature\n",
    "import shapely.geometry as sgeom\n",
    "from datetime import datetime,timedelta\n",
    "from tqdm import tqdm\n",
    "from global_land_mask import globe\n",
    "from typlot.scripts.gsj_typhoon import tydat,see,count_rapidgrow,tydat_CMA,average_datetime,split_str_id,load_land_polygons,detect_landfall\n",
    "from geopy.distance import geodesic\n",
    "import matplotlib.ticker as ticker\n",
    "from typlot.config.global_config import *\n",
    "from typlot.scripts.gsj_typhoon import save_pickle,load_pickle\n",
    "\n",
    "all_ini_time_mode = ['00_12']\n",
    "tynames,tyids = split_str_id(names)\n",
    "draw_obs_opt = True\n",
    "all_obs_baseline= ['RI']  # 'land' 'RI'\n",
    "\n",
    "# ========== 修改 1: 调整 figure 布局，为 colorbar 预留空间 ==========\n",
    "fig, axes = plt.subplots(4, 2, figsize=(15, 20))\n",
    "axes = axes.flatten()\n",
    "\n",
    "# 调整子图位置，底部留出空间\n",
    "fig.subplots_adjust(bottom=0.08)  # 为 colorbar 预留空间\n",
    "\n",
    "# 用于收集所有数据的最大最小值\n",
    "all_data_min = float('inf')\n",
    "all_data_max = float('-inf')\n",
    "all_dis_results = []  # 存储所有距离数据\n",
    "all_draw_data_obs = []  # 存储所有观测数据\n",
    "all_dates_names = []  # 存储所有日期名称\n",
    "\n",
    "for ini_time_mode in all_ini_time_mode:\n",
    "    for obs_baseline in all_obs_baseline:\n",
    "        for axesi,(ty,tyid,year) in enumerate(zip(tynames,tyids,years)):\n",
    "                    \n",
    "            ''' init '''\n",
    "            RIstd = 7\n",
    "            # dt = timedelta(days=1)\n",
    "            dt = timedelta(hours=6)\n",
    "            directory = os.path.join(global_ensdir,f\"{ty}_{tyid}\")\n",
    "            dates_name = sorted(os.listdir(directory))\n",
    "            dates_name = [i for i in dates_name if i[-2:] in ini_time_mode ]\n",
    "            dates_paths = [os.path.join(directory,date) for date in dates_name if date[-2:] in ini_time_mode]\n",
    "            obs_path = os.path.join(global_obsdir,f\"{ty}_CMAobs.txt\")\n",
    "            pic_savepath=os.path.join(global_picdir,f'{obs_baseline}_heatmap',ini_time_mode,'start_day_distance')\n",
    "            os.makedirs(pic_savepath,exist_ok=True)\n",
    "                \n",
    "            if (not os.path.exists(os.path.join(paperdir,f\"fig6_{ty}.pickle\"))) or (not os.path.exists(os.path.join(paperdir,f\"fig6_{ty}_obs.pickle\"))):\n",
    "                \n",
    "                ''' calc 2d numbers '''\n",
    "                # 遍历所有起报时间\n",
    "                max_delta_time = timedelta(days=0)\n",
    "                for j,i in enumerate(dates_paths):\n",
    "                    mem_paths = glob.glob( os.path.join(i,\"TRACK_ID_*\") )\n",
    "                    start_time = datetime.strptime(dates_name[j],\"%Y%m%d%H\")\n",
    "                    # 遍历所有成员\n",
    "                    for single_mem_path in mem_paths:\n",
    "                        single_mem = tydat(single_mem_path,RIstd) \n",
    "                        delta_time = single_mem.time[-1] - start_time\n",
    "                        if delta_time>max_delta_time:\n",
    "                            max_delta_time = delta_time\n",
    "                max_dt_num = math.ceil(max_delta_time/dt) \n",
    "                \n",
    "                ''' 求各起报时间的图 '''\n",
    "                tyobs = tydat_CMA(obs_path)\n",
    "                rcd_obs = count_rapidgrow(RIstd, tyobs.umax, tyobs.time)\n",
    "                num_result = np.zeros(( max_dt_num,len(dates_name) ))\n",
    "                dis_result = np.zeros(( max_dt_num,len(dates_name) ))\n",
    "                dis_var_result = np.zeros(( max_dt_num,len(dates_name) ))\n",
    "                \n",
    "                \n",
    "                for i in tqdm( range(len(dates_paths)), desc=f'processing  {ty} {obs_baseline}_{ini_time_mode}'):\n",
    "                    '''读取这个起报时间的所有集合'''\n",
    "                    mem_paths = glob.glob( os.path.join(dates_paths[i],\"TRACK_ID_*\") )\n",
    "                    sorted_mem_paths = sorted(mem_paths, key=lambda x: int(x.split('_')[-1]))\n",
    "                    ty_mem = [tydat(i,RIstd) for i in sorted_mem_paths]\n",
    "                    start_dt = datetime.strptime( dates_name[i],\"%Y%m%d%H\" )\n",
    "                    for j in range(max_dt_num):\n",
    "                        t0 = j\n",
    "                        t1 = j+1\n",
    "                        t0_dt = start_dt + timedelta(days=t0)\n",
    "                        t1_dt = start_dt + timedelta(days=t1)\n",
    "                        '''如果观测在这两天没有RI，就直接跳过'''\n",
    "                        idx_obs_2d = (tyobs.time>=t0_dt)&(tyobs.time<t1_dt)\n",
    "                        idx_obs_2d_RI = (rcd_obs[idx_obs_2d]==1)\n",
    "                        if np.sum(rcd_obs[idx_obs_2d]) ==0.:\n",
    "                            ''' only record numbers here '''\n",
    "                            for k in range(len(ty_mem)):\n",
    "                                '''cycle through all members : 如果这个成员在这两天没有RI,就跳过'''\n",
    "                                idx_mem_2d    = (ty_mem[k].time>=t0_dt)&(ty_mem[k].time<t1_dt)\n",
    "                                idx_mem_2d_RI = (ty_mem[k].num_rapidgrow()[idx_mem_2d]==1)\n",
    "                                num_result[j,i] += np.sum(idx_mem_2d_RI)\n",
    "                            continue\n",
    "                        else:\n",
    "                            '''如果2d内有多个RI,取平均'''\n",
    "                            time_obs_ave = average_datetime(tyobs.time[idx_obs_2d][idx_obs_2d_RI])\n",
    "                            lon_obs_ave  = np.mean(tyobs.lon[idx_obs_2d][idx_obs_2d_RI])\n",
    "                            lat_obs_ave  = np.mean(tyobs.lat[idx_obs_2d][idx_obs_2d_RI])\n",
    "                            '''开始计算'''\n",
    "                            num = 0 \n",
    "                            dis_list = [] \n",
    "            \n",
    "                            for k in range(len(ty_mem)):\n",
    "                                '''cycle through all members : 如果这个成员在这两天没有RI,就跳过'''\n",
    "                                idx_mem_2d    = (ty_mem[k].time>=t0_dt)&(ty_mem[k].time<t1_dt)\n",
    "                                idx_mem_2d_RI = (ty_mem[k].num_rapidgrow()[idx_mem_2d]==1)\n",
    "                                time_mem = ty_mem[k].time[idx_mem_2d][idx_mem_2d_RI]\n",
    "                                lat_mem  = ty_mem[k].lat[idx_mem_2d][idx_mem_2d_RI]\n",
    "                                lon_mem  = ty_mem[k].lon[idx_mem_2d][idx_mem_2d_RI]\n",
    "                                if np.sum(idx_mem_2d_RI) ==0.:\n",
    "                                    continue\n",
    "                                else:\n",
    "                                    num+=np.sum(idx_mem_2d_RI)\n",
    "                                    for l in range(len(lon_mem)):\n",
    "                                        dis_list.append(geodesic((lat_obs_ave,lon_obs_ave),(lat_mem[l],lon_mem[l])).kilometers)\n",
    "                        if num==0:\n",
    "                            continue\n",
    "                        num_result[j,i] = num\n",
    "                        dis_result[j,i] = np.array(dis_list).mean()\n",
    "                        dis_var_result[j,i] = np.array(dis_list).std()\n",
    "                        \n",
    "                if draw_obs_opt == True : \n",
    "                    ''' init obsCMA; RI points reord '''\n",
    "                    obs = tydat_CMA(obs_path) \n",
    "                    draw_data_obs = np.zeros( (len(dates_name),max_dt_num),dtype=int )\n",
    "                    tt,dd = draw_data_obs.shape\n",
    "                    \n",
    "                    if obs_baseline == 'RI':\n",
    "                        record_obs = count_rapidgrow(RIstd,obs.umax, obs.time)\n",
    "                        for t in range(tt):\n",
    "                            #预报时效\n",
    "                            for d in range(dd):\n",
    "                                try:\n",
    "                                    start_time = datetime.strptime(dates_name[t],\"%Y%m%d%H\")\n",
    "                                    cal_time = start_time + d*timedelta(days=1)\n",
    "                                    calend_time = cal_time + timedelta(days=1)\n",
    "                                    # 对齐\n",
    "                                    trange = (obs.time>=cal_time) & (obs.time<calend_time)\n",
    "                                    draw_data_obs[t][d] = 1 if np.sum(record_obs[trange])>0 else 0 \n",
    "                                except Exception as e:\n",
    "                                    print('line 132',e)\n",
    "\n",
    "                elif obs_baseline=='land':\n",
    "                    # get  land  time\n",
    "                    land_polys=load_land_polygons(os.path.join(global_shpdir,'China','bou1_4p.shp'),os.path.join(global_shpdir,'China','bou1_4p.dbf'))\n",
    "                    flags,id1 = detect_landfall(obs.lat, obs.lon, land_polys)\n",
    "                    land_time = obs.time[id1]\n",
    "                    # identify  position\n",
    "                    for t in range(tt):\n",
    "                        for d in range(dd):\n",
    "                            try:\n",
    "                                start_time = datetime.strptime(dates_name[t],\"%Y%m%d%H\")\n",
    "                                cal_time = start_time + d*timedelta(days=1)\n",
    "                                calend_time = cal_time + timedelta(days=1)\n",
    "                                trange = (obs.time>=cal_time) & (obs.time<calend_time)\n",
    "                                draw_data_obs[t][d] = 1 if land_time in obs.time[trange] else 0\n",
    "                            except Exception as e:\n",
    "                                print('line 149',e)\n",
    "                \n",
    "                # 保存数据\n",
    "                all_dis_results.append(dis_result)\n",
    "                all_draw_data_obs.append(draw_data_obs)\n",
    "                all_dates_names.append(dates_name)\n",
    "                \n",
    "                # 收集数据范围（注意：dis_result 的维度是 [时效, 起报时间]）\n",
    "                # 需要转置来匹配显示格式 [起报时间, 时效]\n",
    "                var_display = dis_result[:15,:].T  # 转置，取前15天\n",
    "                # 忽略0值，只统计有效值\n",
    "                valid_data = var_display[var_display > 0]\n",
    "                if len(valid_data) > 0:\n",
    "                    all_data_min = min(all_data_min, valid_data.min())\n",
    "                    all_data_max = max(all_data_max, valid_data.max())\n",
    "                \n",
    "                save_pickle(draw_data_obs,os.path.join(paperdir,f\"fig6_{ty}.pickle\"))\n",
    "                save_pickle(dis_result,os.path.join(paperdir,f\"fig6_{ty}_obs.pickle\"))\n",
    "            else: \n",
    "                draw_data_obs = load_pickle(os.path.join(paperdir,f\"fig6_{ty}.pickle\"))\n",
    "                dis_result = load_pickle(os.path.join(paperdir,f\"fig6_{ty}_obs.pickle\"))\n",
    "\n",
    "            ''' 绘制热图 '''\n",
    "            ax = axes[axesi]\n",
    "            \n",
    "            # ========== 修改 2: 使用 imshow 并关闭每个子图的 colorbar ==========\n",
    "            # 11.19 将0修改为np.nan, 来实现换成白色\n",
    "            dis_result[dis_result==0] = np.nan\n",
    "\n",
    "            # 注意：dis_result 维度是 [时效, 起报时间]，需要转置显示为 [起报时间, 时效]\n",
    "            im = ax.imshow(dis_result[:15,:].T, cmap=\"YlGnBu\", aspect='auto', \n",
    "                          vmin=all_data_min, vmax=all_data_max)\n",
    "            \n",
    "            # 手动添加 heatmap 的其他元素\n",
    "            # 添加网格线\n",
    "            for i in range(len(dates_name) + 1):\n",
    "                ax.axhline(i - 0.5, color='white', linewidth=0.5)\n",
    "            for j in range(15 + 1):\n",
    "                ax.axvline(j - 0.5, color='white', linewidth=0.5)\n",
    "            \n",
    "            # 设置刻度\n",
    "            ax.set_yticks(np.arange(len(dates_name))[::4]) \n",
    "            yticklabels = [i[4:] for i in dates_name[::4]]\n",
    "            ax.set_yticklabels(yticklabels, rotation=0) \n",
    "            \n",
    "            ax.set_xticks(np.arange(15))\n",
    "            xlabels=[str(i)+\"d\"  for i in range(1,16)]\n",
    "            ax.set_xticklabels(xlabels, rotation=0)\n",
    "\n",
    "            # 添加台风名称到左下\n",
    "            ax.text(\n",
    "                0.02, 0.03,\n",
    "                f\"{ty}({year})\",\n",
    "                transform=ax.transAxes,\n",
    "                fontsize=10,\n",
    "                fontweight='bold',\n",
    "                color='white',\n",
    "                verticalalignment='bottom',\n",
    "                horizontalalignment='left',\n",
    "                bbox=dict(\n",
    "                    boxstyle='round,pad=0.5',\n",
    "                    edgecolor='none',\n",
    "                    alpha=0.7\n",
    "                )\n",
    "            )\n",
    "            \n",
    "            ''' plot scatter points '''\n",
    "            # draw_data_obs 维度是 [起报时间, 时效]，找出值为1的位置\n",
    "            points = np.where(draw_data_obs[:,:15] == 1)\n",
    "            x_coords = points[1]  # 时效维度（列）\n",
    "            y_coords = points[0]  # 起报时间维度（行）\n",
    "            ax.scatter(x_coords, y_coords, color='red', marker='o',s=10)\n",
    "            ax.invert_yaxis()\n",
    "\n",
    "# ========== 修改 3: 添加统一的 colorbar ==========\n",
    "# 在底部添加 colorbar\n",
    "cbar_ax = fig.add_axes([0.15, 0.03, 0.7, 0.015])  # [left, bottom, width, height]\n",
    "cbar = fig.colorbar(im, cax=cbar_ax, orientation='horizontal')\n",
    "cbar.set_label('MAE (km)', fontsize=12)\n",
    "cbar.locator = ticker.MaxNLocator(integer=True)\n",
    "cbar.update_ticks()\n",
    "\n",
    "fig.suptitle(\"Mean Distance Between Ensemble RI and Observed RI Locations\",\n",
    "              fontsize=18, fontweight='bold', y=0.99)\n",
    "fig.subplots_adjust(left=0.1, right=0.9, top=0.95, bottom=0.08, wspace=0.3, hspace=0.3)  # wspace 和 hspace 控制子图间隔\n",
    "\n",
    "# 在循环结束后添加标签\n",
    "labels = [\"(a)\",\"(b)\",\"(c)\",\"(d)\",\"(e)\",\"(f)\",\"(g)\",\"(h)\"]\n",
    "for ax, label in zip(axes.flatten(), labels):\n",
    "    ax.text(0.9, 0.95, label,\n",
    "            transform=ax.transAxes,  # 使用轴坐标系 (0-1)\n",
    "            fontsize=14,\n",
    "            fontweight='bold',\n",
    "            va='top',                # 垂直对齐：顶部\n",
    "            ha='left',              # 水平对齐：右侧\n",
    "            bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))\n",
    "\n",
    "if draw_opt==False:\n",
    "    plt.tight_layout(rect=[0, 0.05, 1, 0.99])  # 为 colorbar 和 title 留空间\n",
    "    plt.savefig(os.path.join(paperdir, f\"fig6.{pictype}\"), dpi=900)\n",
    "    plt.close()\n",
    "else:\n",
    "    plt.show()\n",
    "\n",
    "plt.clf()\n",
    "plt.close('all')\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "wmq",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
